﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Diagnostics;

namespace MSAS.DIP.Common
{
    public class RetryUtility
    {
        #region Fields

        private RetryWait m_waitSchedule;

        #endregion

        #region Properties

        /// <summary>
        /// Called before every retry.
        /// </summary>
        public event EventHandler<RetryArgs> RetryHandler;

        /// <summary>
        /// Called upon final failure
        /// </summary>
        public event EventHandler<RetryArgs> FailureHandler;

        /// <summary>
        /// Called upon successful try.
        /// </summary>
        public event EventHandler<RetryArgs> SuccessHandler;

        /// <summary>
        /// Default retry schedule for work units
        /// </summary>
        public RetryWait WaitSchedule 
        {
            get
            {
                return m_waitSchedule;
            }

            set
            {
                if (value.InitialWaitTime <= TimeSpan.Zero)
                {
                    throw new ArgumentOutOfRangeException("value.InitialWaitTime", "The initial wait must be positive");
                }
                if (value.WaitScalingFactor <= 1)
                {
                    throw new ArgumentOutOfRangeException("value.WaitScalingFactor", "The scaling factor must be greater than 1");
                }
                if (value.TotalWaitTime <= TimeSpan.Zero)
                {
                    throw new ArgumentOutOfRangeException("value.TotalWaitTime", "The total wait must be positive");
                }

                m_waitSchedule = value;
            }
        }

        #endregion

        #region Constructors

        public RetryUtility()
        {
            // Set up the default wait.
            WaitSchedule = RetryWait.ShortWait;
        }

        public RetryUtility(RetryWait waitSchedule)
        {
            WaitSchedule = waitSchedule;
        }

        #endregion

        #region Helpers

        private IEnumerable<TimeSpan> UnrolledWaitSchedule(
            TimeSpan initialWaitTime, 
            double waitScalingFactor, 
            TimeSpan totalWaitTime)
        {
            // Scale the times starting at the initial wait time
            TimeSpan totalTimeSoFar = TimeSpan.Zero;
            TimeSpan nextWait = TimeSpan.Zero;

            for (double i = initialWaitTime.Ticks; 
                 totalTimeSoFar <= totalWaitTime; 
                 i *= waitScalingFactor)
            {
                nextWait = TimeSpan.FromTicks((long)i);
                TimeSpan nextTotalValue = totalTimeSoFar.Add(nextWait);
                if (nextTotalValue > totalWaitTime)
                {
                    break;
                }
                yield return nextWait;
                totalTimeSoFar = nextTotalValue;
            }

            // If the scaling resulted in a remainder then try
            // after waiting the remainder
            TimeSpan remainder = totalWaitTime.Subtract(totalTimeSoFar);
            if (remainder > TimeSpan.Zero)
            {
                yield return remainder;
            }

            // Always make one last try at the end
            yield return TimeSpan.Zero;
        }

        private void IndicateSuccess(object sender, RetryArgs e)
        {
            if (SuccessHandler != null)
            {
                SuccessHandler(sender, e);
            }
        }

        private void IndicateRetry(object sender, RetryArgs e)
        {
            if (RetryHandler != null)
            {
                RetryHandler(sender, e);
            }
        }

        private void IndicateFailure(object sender, RetryArgs e)
        {
            if (FailureHandler != null)
            {
                FailureHandler(sender, e);
            }
        }

        #endregion

        #region Private Methods

        private T TryImplementation<T>(
            TimeSpan initialWaitTime,
            double waitScalingFactor,
            TimeSpan totalWaitTime,
            Func<T> workToRetry)
        {
            Debug.Assert(initialWaitTime > TimeSpan.Zero, "The initial wait time should be positive");
            Debug.Assert(waitScalingFactor > 1.0, "The scaling factor should be greater than 1.0");
            Debug.Assert(totalWaitTime > TimeSpan.Zero, "The total wait time should be positive");

            int attemptCounter = 0;
            RetryArgs retryStatus = new RetryArgs(1, TimeSpan.Zero, null);

            foreach (TimeSpan nextSleep in UnrolledWaitSchedule(initialWaitTime, waitScalingFactor, totalWaitTime))
            {
                attemptCounter++;
                try
                {
                    T returnValue = workToRetry();
                    IndicateSuccess(workToRetry.Target, retryStatus);
                    return returnValue;
                }
                catch (Exception ex)
                {
                    retryStatus = new RetryArgs(attemptCounter, nextSleep, ex);
                    if (nextSleep != TimeSpan.Zero)
                    {
                        // Indicate a retry and sleep
                        IndicateRetry(workToRetry.Target, retryStatus);
                        if (retryStatus.Cancel)
                        {
                            break;
                        }

                        Thread.Sleep(nextSleep);
                    }
                }
            }

            // If we got here, then all attempts have been exhausted or we got cancelled
            IndicateFailure(workToRetry.Target, retryStatus);
            throw retryStatus.ThrownException;
        }

        #endregion

        #region Public Interface

        /// <summary>
        /// Try a function and specify the wait schedule
        /// </summary>
        public T Try<T>(TimeSpan initialWait, double waitScalingFactor, TimeSpan totalWaitTime, Func<T> workToRetry)
        {
            if (initialWait <= TimeSpan.Zero)
            {
                throw new ArgumentOutOfRangeException("initialWait", "The initial wait must be positive");
            }
            if (waitScalingFactor <= 1)
            {
                throw new ArgumentOutOfRangeException("waitScalingFactor", "The scaling factor must be greater than 1");
            }
            if (totalWaitTime <= TimeSpan.Zero)
            {
                throw new ArgumentOutOfRangeException("totalWaitTime", "The total wait must be positive");
            }

            return TryImplementation(initialWait, waitScalingFactor, totalWaitTime, workToRetry);
        }

        /// <summary>
        /// Try a function and specify the wait schedule
        /// </summary>
        public void Try(TimeSpan initialWait, double waitScalingFactor, TimeSpan totalWaitTime, Action workToRetry)
        {
            // Try the action as a function that returns null
            Try<object>(initialWait, waitScalingFactor, totalWaitTime, () => { workToRetry(); return null; });
        }

        /// <summary>
        /// Try a function with a pre-set wait schedule
        /// </summary>
        public T Try<T>(Func<T> workToRetry)
        {
            return TryImplementation(WaitSchedule.InitialWaitTime, WaitSchedule.WaitScalingFactor, WaitSchedule.TotalWaitTime, workToRetry);
        }

        /// <summary>
        /// Try an action with a pre-set wait schedule
        /// </summary>
        public void Try(Action workToRetry)
        {
            TryImplementation<object>(WaitSchedule.InitialWaitTime, WaitSchedule.WaitScalingFactor, WaitSchedule.TotalWaitTime, () => { workToRetry(); return null; });
        }

        #endregion
    }
}
